import time

import numpy as np
import torch
from torch import optim
from torch.nn import functional as F
from sklearn.metrics import accuracy_score, classification_report, confusion_matrix

from utils import get_time_dif


def evaluate(config, model, dev_iter, test=False):
    model.eval()
    loss_total = 0.0
    predict_all = []
    label_all = []
    with torch.no_grad():
        for X, y, _ in dev_iter:
            outputs = model(X.to(config.device))
            loss = F.cross_entropy(outputs, y.to(config.device))
            loss_total += loss
            label = y.data.cpu()
            predict = torch.max(outputs, dim=1)[1].cpu()
            label_all = np.append(label_all, label)
            predict_all = np.append(predict_all, predict)

        acc = accuracy_score(label_all, predict_all)

    if test:
        report = classification_report(label_all, predict_all, target_names=config.class_list, digits=4)
        confusion = confusion_matrix(label_all, predict_all)
        return acc, loss_total / len(dev_iter), report, confusion
    return acc, loss_total / len(dev_iter)


def test(config, model, test_iter):
    model.load_state_dict(torch.load(config.save_path))
    model.eval()
    start_time = time.time()
    test_acc, test_loss, test_report, test_confusion = evaluate(config, model, test_iter, True)
    msg = "Test Loss:{0:>6.2}, Test Acc:{1:>6.2%}, time:{2:>6}s"
    print(msg.format(test_loss, test_acc, get_time_dif(start_time)))
    print("Precision, recall, F1-score")
    print(test_report)
    print("Confusion Matrix")
    print(test_confusion)



def train(config, model, train_iter, dev_iter, test_iter, writer):
    print("Start Training...")
    start_train_time = time.time()
    model.train()
    optimizer = optim.Adam(model.parameters(), lr=config.learning_rate)

    total_batch = 0
    dev_best_loss = float('inf')
    last_improvement = 0
    flag = False
    for epoch in range(config.n_epochs):
        print('Epoch{}/{}'.format(epoch + 1, config.n_epochs))
        epoch_time = time.time()

        for i, (X, y, length) in enumerate(train_iter):
            model.train()
            outputs = model(X.to(config.device))
            model.zero_grad()
            loss = F.cross_entropy(outputs, y.to(config.device))
            loss.backward()
            optimizer.step()

            if total_batch % 1000 == 0:
                true = y.data.cpu()
                predict = torch.max(outputs.data, dim=1)[1].cpu()

                train_acc = accuracy_score(true, predict)
                dev_acc, dev_loss = evaluate(config, model, dev_iter)
                if dev_loss < dev_best_loss:
                    dev_best_loss = dev_loss
                    torch.save(model.state_dict(), config.save_path)
                    improve = "*"
                    last_improvement = total_batch
                else:
                    improve = ""

                time_epoch_diff = get_time_dif(start_train_time)
                msg = 'Epoch: {0:>6}, Iter: {1:>6}, Train Loss: {2:>5.2f}, Train Acc: {3:>6.2%}, ' \
                      'Val Loss: {4:>5.2f},Val Acc: {5:>6.2%}, Time: {6:.2f}s'
                print(msg.format(epoch+1, total_batch, loss.item(), train_acc, dev_loss, dev_acc, time_epoch_diff))
                writer.add_scalar("loss/train", loss.item(), total_batch)
                writer.add_scalar("loss/dev", dev_loss, total_batch)
                writer.add_scalar("acc/train", train_acc, total_batch)
                writer.add_scalar("acc/dev", dev_acc, total_batch)

            total_batch += 1
            if total_batch - last_improvement > config.require_improvement:
                print("没有优化啦！")
                flag = True
                break

        if flag:
            break

    writer.close()
    test(config, model, test_iter)
